import gradio as gr
import pickle
import cv2
import numpy as np
from PIL import Image

# 加载最佳模型
with open('best_model.pkl', 'rb') as f:
    model = pickle.load(f)

def predict_image(image):
    # 将 PIL.Image.Image 对象转换为 numpy 数组
    if isinstance(image, Image.Image):
        image = np.array(image)
    
    # 如果是灰度图，转换为 RGB
    if len(image.shape) == 2:  # 如果是灰度图
        img = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    else:  # 如果已经是 RGB 图
        img = image
        
    # 调整尺寸为 64x64
    img = cv2.resize(img, (64, 64))
    
    # 确保图像是 64x64x3 的形状
    if img.shape[0] == 64 and img.shape[1] == 64 and img.shape[2] == 3:
        # 扁平化处理，适应模型输入
        image_flattened = img.reshape(1, -1)  # 展平图像
    else:
        raise ValueError("图像预处理后的形状不正确，应为64x64x3。")
    
    # 随机抽样到4096个特征
    sample_indices = np.random.choice(image_flattened.shape[1], 4096, replace=False)
    image_flattened_sampled = image_flattened[:, sample_indices]
    
    # 模型预测
    prediction = model.predict(image_flattened_sampled)[0]
    return {
        "猫": float(prediction == 0),
        "狗": float(prediction == 1)
    }

# 创建Gradio接口
iface = gr.Interface(
    fn=predict_image, 
    inputs=gr.components.Image(type="pil", tool="editor"),  # 使用 gradio.components
    outputs=gr.components.Label(num_top_classes=2),  # 使用 gradio.components
    title="猫狗分类器"
)

# 启动Web应用
iface.launch()